import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.preprocessing import StandardScaler
from keras.utils.np_utils import to_categorical
from keras.models import Sequential
from keras.layers import Dense, LSTM
from sklearn.metrics import classification_report, accuracy_score

# Import dataset
raw_data = pd.read_csv('./datingTestSet2.txt', delimiter='\t', header=None)
print("Dataset Shape: " + str(raw_data.shape))

# Preprocess data
raw_data = raw_data.drop_duplicates()
raw_data = raw_data.fillna(raw_data.mean())
# Normalize data using z-score
ss = StandardScaler()
# 0: flight, 1: ice cream, 2: game time
scale_features = [0, 1, 2]
raw_data[scale_features] = ss.fit_transform(raw_data[scale_features])
raw_data.hist(grid=False, figsize=(12, 12))
plt.show()

# Separate data
raw_data = np.array(raw_data)
data = raw_data[:, :2]
target = raw_data[:, 3]
print("data shape")
print(data.shape)
print("target shape")
print(target.shape)
for i in range(len(target)):
    target[i] = target[i] - 1
# print(target)
# Convert to one-hot
plain_target = target
target = to_categorical(target, 3)
# print(target[0:3, :])
x, y = data, target

# Convert into LSTM input
train_x, train_y = x[:500, :], y[:500]
test_x, test_y = x[500:, :], y[500:]
# Reshape into 3D [samples, timesteps, features]
train_x = train_x.reshape((train_x.shape[0], 1, train_x.shape[1]))
test_x = test_x.reshape((test_x.shape[0], 1, test_x.shape[1]))
print(train_x.shape, train_y.shape, test_x.shape, test_y.shape)

# Model setup
model = Sequential()
model.add(LSTM(50, input_shape=(train_x.shape[1], train_x.shape[2])))
# model.add(Dense(units=8, input_dim=2, activation='relu'))
model.add(Dense(units=3, activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(train_x, train_y, epochs=150, batch_size=10, validation_data=(test_x, test_y), verbose=2, shuffle=False)

# Evaluate model
result = model.evaluate(test_x, test_y)
print(result[1])

# Predict model
proba = model.predict(test_x)
print(proba[0:3])
classes = model.predict_classes(test_x)
print(classes[0:3])

# Print classification report
predict_result = []
for i in proba:
    max_index, max_proba = 0, 0
    for j in range(3):
        if i[j] >= max_proba:
            max_index = j
            max_proba = i[j]
    predict_result.append(max_index)
target_names = ['class 1', 'class 2', 'class 3']
print(classification_report(plain_target[500:], predict_result, target_names=target_names))
